#!/usr/bin/python
'''
Takes an AGP file and rebuilds it given contig sizes.
'''
import sys

# parameters.
agp_file = sys.argv[1]
size_file = sys.argv[2]
trans_file = sys.argv[3]

# read in translate.
#print "read trans"
trans = {}
fin = open(trans_file, "rb")
for line in fin:
	tmp = line.strip().split()
	trans[tmp[0]] = tmp[1]
fin.close()

# read in  contig sizes.
#print "read size"
ctgsz = {}
fin = open(size_file, "rb")
for line in fin:
	tmp = line.strip().split()
	ctgsz[tmp[1]] = int(tmp[0])
fin.close()

# read in agp, removing gaps.
#print "load scaffold"
scafs = {}
gaps = {}
fin = open(agp_file, "rb")
i = 0
for line in fin:
	# split line.
	tmp = line.strip().split("\t")
	
	# check if contig isn't used anymore.
	if tmp[4] == "W":
		if tmp[5] not in trans: continue
		tmp[5] = trans[tmp[5]]
	
	
	# check if scaffold exists.
	if tmp[0] not in scafs:
		scafs[tmp[0]] = []
	
	# skip non-contigs.
	if tmp[4] == "W" and tmp:
		# add contig info.
		scafs[tmp[0]].append(tmp)

	# debug.
	i += 1
	#if i > 2000: break
fin.close()

# rebuild the AGP.
#print "rebuild scaffold"
for sid in scafs:
	# initialize idx.
	idx = 1
	cidx = 1
	numc = len(scafs[sid]) - 1
	
	# build each row.
	for row in scafs[sid]:
		# Add gap if not first.
		if cidx != 1:
			# build gap.
			gstart = idx
			gstop = int(row[1]) - 1 
			
			dat = [sid, str(gstart), str(gstop), str(cidx), "N", str(gstop-gstart+1), "fragment", "yes",""]
			
			# write gap.
			print '\t'.join(dat)
			
			# update stuff.
			idx = gstop + 1
			cidx += 1
		
		# Build row.
		ctg = row[5]
		start = idx
		stop = start + ctgsz[ctg] - 1
		
		cstart = 1
		cstop = cstart + ctgsz[ctg] - 1
		
		dat = [sid, str(start), str(stop), str(cidx), "W", ctg, str(cstart), str(cstop), "+"]
		
		# write row.
		print '\t'.join(dat)

		# update stuff.
		idx = stop + 1
		cidx += 1
		
